# -*- coding: utf-8 -*-
"""
WebRTC连接管理类
负责管理WebRTC连接、数据通道和事件处理
"""

import asyncio
from aiortc import RTCPeerConnection, RTCSessionDescription, RTCIceCandidate
from common.utils import LogUtil
from dev_server.device_context import DeviceContext
from dev_server.webrtc_video_track import WebRTCVideoTrack

logger = LogUtil.get_logger(__name__)

class WebRTCConnection:
    """WebRTC连接管理"""
    
    def __init__(self, client_id: str, device_sn: str, device_context: DeviceContext):
        self.client_id = client_id
        self.device_sn = device_sn
        self.device_context = device_context
        self.pc = None
        self.candidate_cache = []
        self.video_track = None
        self.uitree_channel = None
        self.uinput_channel = None
        self.perf_channel = None
        self.ctrl_channel = None
        self.running = False
        self.pc = RTCPeerConnection()

    def get_ice_candidate(self, seq: int):
        """获取ICE候选"""
        if len(self.candidate_cache) > seq:
            return self.candidate_cache[seq]
        return None

    async def close(self):
        """关闭WebRTC连接"""
        if self.video_track:
            self.video_track.stop()
            self.video_track = None
        self.running = False
        if self.pc:
            await self.pc.close()
            self.pc = None

    async def setup_pc_handlers(self):
        """设置PeerConnection事件处理器"""
        pc = self.pc

        @pc.on("icecandidate")
        def on_icecandidate(candidate):
            if candidate:
                self.candidate_cache.append(candidate)
        
        @pc.on("datachannel")
        def on_datachannel(channel):
            logger.info(f"建立WebRTC数据通道: {channel.label}")

            @channel.on("close")
            def on_close():
                logger.info(f"关闭数据通道: {channel.label}")
            
            if channel.label == 'uitree':
                self.uitree_channel = channel
                self.uinput_channel = channel
                @channel.on("message")
                def on_message(message):
                    logger.info(f"收到WebRTC数据: channel=uinput, message={message}")
                    
            if channel.label == 'perf':
                self.perf_channel = channel
                self.ctrl_channel = channel
                @channel.on("message")
                def on_message(message):
                    logger.info(f"收到WebRTC数据: channel=ctrl, message={message}")

            if channel.label != 'uitree' and channel.label != 'perf':
                logger.error(f"未知数据通道: {channel.label}")

        @pc.on("connectionstatechange")
        async def on_connectionstatechange():
            state = self.pc.connectionState
            logger.info(f"WebRTC连接状态变化: {state}")
            
            if state == 'failed' or state == 'disconnected':
                logger.info(f"WebRTC连接已断开: {state}")
                await self.close()
            elif state == 'connected':
                asyncio.create_task(self.__start_screen_transmission())
                asyncio.create_task(self.__start_uitree_transmission())
                asyncio.create_task(self.__start_uinput_transmission())
                asyncio.create_task(self.__start_perf_transmission())
                asyncio.create_task(self.__start_ctrl_transmission())

    async def set_remote_description(self, offer) -> bool:
        try:
            logger.info(f"设置远程描述: type={offer['type']}, sdp长度={len(offer['sdp'])}")
            await self.pc.setRemoteDescription(RTCSessionDescription(
                sdp=offer['sdp'],
                type=offer['type']
            ))
            logger.info("远程描述设置成功")
            return True
        except Exception as desc_error:
            logger.error(f"设置远程描述失败: {desc_error}")
            return False

    def set_video_track(self) -> bool:
        try:
            self.video_track = WebRTCVideoTrack()
            self.pc.addTrack(self.video_track)
            logger.info("视频轨道设置成功")
        except Exception as e:
            logger.error(f"设置视频轨道失败: {e}")
            return False

        try:
            # 确保所有transceiver都有正确的方向设置
            for transceiver in self.pc.getTransceivers():
                if transceiver.sender.track and hasattr(transceiver.sender.track, 'kind'):
                    if transceiver.sender.track.kind == 'video':
                        # 远程服务器发送视频，所以设置为sendonly
                        transceiver.direction = 'sendonly'
                        logger.info(f"设置发送视频轨道方向: {transceiver.direction}")
                elif hasattr(transceiver, 'receiver') and transceiver.receiver.track:
                    # 对于接收轨道，设置为recvonly
                    if hasattr(transceiver.receiver.track, 'kind') and transceiver.receiver.track.kind == 'video':
                        transceiver.direction = 'recvonly'
                        logger.info(f"设置接收视频轨道方向: {transceiver.direction}")
        except Exception as direction_error:
            logger.error(f"设置轨道方向失败: {direction_error}")
            return False
        
        return True
    
    async def create_answer(self):
        """创建WebRTC answer"""
        try:
            logger.info("开始创建answer")
            answer = await self.pc.createAnswer()
            logger.info(f"answer创建成功: type={answer.type}, sdp长度={len(answer.sdp)}")
            logger.info(f"answer SDP内容: {answer.sdp[:200]}...")  # 打印SDP前200字符用于调试
            
            logger.info("开始设置本地描述")
            await self.pc.setLocalDescription(answer)
            logger.info("本地描述设置成功")
            return {
                    'sdp': self.pc.localDescription.sdp,
                    'type': self.pc.localDescription.type
                }
        except Exception as answer_error:
            logger.error(f"创建answer失败: {answer_error}")
            import traceback
            logger.error(f"详细错误信息: {traceback.format_exc()}")
            return None
    
    async def add_ice_candidate(self, candidate_data):
        """添加ICE候选"""
        try:
            # 处理前端发送的RTCIceCandidate对象格式
            if not isinstance(candidate_data, dict):
                logger.error(f"无效的ICE候选格式")
                return False
            if 'candidate' in candidate_data and 'sdpMid' in candidate_data:
                # 标准WebRTC格式 - 解析candidate字符串
                candidate_str = candidate_data['candidate']
                sdp_mid = candidate_data['sdpMid']
                sdp_mline_index = candidate_data.get('sdpMLineIndex')
                # 解析candidate字符串
                parts = candidate_str.split()

                if len(parts) < 8:
                    logger.error(f"无法解析ICE候选字符串")
                    return False
                
                foundation = parts[0].split(':')[1]
                component = int(parts[1])
                protocol = parts[2]
                priority = int(parts[3])
                ip = parts[4]
                port = int(parts[5])
                typ = parts[7]
                
                candidate = RTCIceCandidate(
                    component=component,
                    foundation=foundation,
                    ip=ip,
                    port=port,
                    priority=priority,
                    protocol=protocol,
                    type=typ,
                    sdpMid=sdp_mid,
                    sdpMLineIndex=sdp_mline_index
                )
            else:
                # 自定义格式
                candidate = RTCIceCandidate(
                    component=candidate_data.get('component', 1),
                    foundation=candidate_data.get('foundation', ''),
                    ip=candidate_data.get('address', candidate_data.get('ip', '')),
                    port=candidate_data.get('port', 0),
                    priority=candidate_data.get('priority', 0),
                    protocol=candidate_data.get('protocol', 'udp'),
                    type=candidate_data.get('type', 'host')
                )
        
            await self.pc.addIceCandidate(candidate)
            logger.info(f"添加ICE候选成功")
            return True
        except Exception as e:
            logger.error(f"处理ICE候选错误: {e}")
            return False
    
    async def __start_screen_transmission(self):
        pass

    async def __start_uitree_transmission(self):
        """UI树数据传输循环"""
        pass

    async def __start_uinput_transmission(self):
        pass

    async def __start_perf_transmission(self):
        pass

    async def __start_ctrl_transmission(self):
        pass